{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/mlop-ai/mlop/blob/main/examples/transformers.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n",
    "\n",
    "<h1 align=\"center\" style=\"font-family: Inter, sans-serif; font-style: normal; font-weight: 700; font-size: 72px\">m:lop</h1>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -Uq \"mlop[full]\" datasets evaluate accelerate scikit-learn transformers\n",
    "# %pip install \"mlop[full] @ git+https://github.com/mlop-ai/mlop.git\"\n",
    "# import sys; import os; sys.path.insert(0, os.path.dirname(os.path.abspath(os.path.dirname(\"__file__\"))))\n",
    "import mlop\n",
    "from mlop.compat.transformers import MLOPCallback\n",
    "\n",
    "mlop.login()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qbd4wDpXMf5b"
   },
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "import evaluate\n",
    "import numpy as np\n",
    "from transformers import AutoTokenizer, AutoModelForSequenceClassification\n",
    "\n",
    "dataset = load_dataset(\"yahoo_answers_topics\")\n",
    "dataset = dataset.rename_column(\"topic\", \"labels\")\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"distilbert-base-uncased\")\n",
    "dataset = dataset.map(\n",
    "    lambda x: tokenizer(x[\"question_title\"], truncation=True), batched=True\n",
    ")\n",
    "model = AutoModelForSequenceClassification.from_pretrained(\n",
    "    \"distilbert-base-uncased\", num_labels=len(dataset[\"train\"].unique(\"labels\"))\n",
    ")\n",
    "accuracy_metric = evaluate.load(\"accuracy\")\n",
    "\n",
    "\n",
    "def compute_metrics(eval_pred):\n",
    "    predictions, labels = eval_pred\n",
    "    predictions = np.argmax(predictions, axis=1)\n",
    "    return accuracy_metric.compute(predictions=predictions, references=labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import Trainer, TrainingArguments\n",
    "\n",
    "args = TrainingArguments(\n",
    "    report_to=\"mlop\",\n",
    "    output_dir=\".mlop\",\n",
    "    overwrite_output_dir=True,\n",
    "    eval_strategy=\"steps\",\n",
    "    learning_rate=5e-5,\n",
    "    max_steps=100,\n",
    "    logging_steps=10,\n",
    "    eval_steps=30,\n",
    "    save_steps=90,\n",
    "    load_best_model_at_end=True,\n",
    "    metric_for_best_model=\"accuracy\",\n",
    "    run_name=\"custom_training\",\n",
    ")\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=args,\n",
    "    train_dataset=dataset[\"train\"],\n",
    "    eval_dataset=dataset[\"test\"],\n",
    "    tokenizer=tokenizer,\n",
    "    compute_metrics=compute_metrics,\n",
    ")\n",
    "\n",
    "trainer.train()\n",
    "mlop.finish()"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
